{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "[PRIME, ICLR 2022] Colab Parsing Dataset",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "Copyright 2022 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ],
      "metadata": {
        "id": "4rSCDmvABaV4"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "id": "N6ZcTf1hBQHO"
      },
      "outputs": [],
      "source": [
        "#@title License\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#  Instructions to use PRIME dataset parser\n",
        "\n",
        "This colab provides the details to parse the\n",
        "[PRIME](https://arxiv.org/abs/2110.11346) dataset available on Google Cloud Storage: [gs://gresearch/prime](https://console.cloud.google.com/storage/browser/gresearch/prime)\n",
        "\n",
        "You may download the dataset either by using the Google Cloud Storage web interface or using gsutil:\n",
        "\n",
        "```\n",
        "gsutil cp -r gs://gresearch/prime /tmp/prime/\n",
        "```\n",
        "\n",
        "The PRIME dataset contains **10** microarchitectural parameters, including number of compute units, on-chip memory size, and number of SIMD units, along with the Runtime (millisecond) and Chip Area Usage (mm$^2$) for all the nine applications that we evaluated in our paper.\n",
        "\n",
        "The data in this dataset are collected using an industry-grade cycle-accurate simulator. \n",
        "\n",
        "This dataset contains both infeasible and feasible data points as described in\n",
        "[PRIME](https://arxiv.org/abs/2110.11346). The descriptors of the collected\n",
        "data are presented in the table below (Table 1).\n",
        "\n",
        "|                  | # of Infeasible | # of Feasible | Max Runtime (ms) | Min Runtime (ms) | Average Runtime (ms) |\n",
        "|------------------|-----------------|---------------|------------------|------------------|----------------------|\n",
        "| **MobileNetEdgeTPU** |          384355 |        115711 |         16352.26 |           252.22 |               529.13 |\n",
        "| **MobilenetV2**      |          744718 |        255414 |          7398.13 |           191.35 |               375.05 |\n",
        "| **MobilenetV3**      |          797460 |        202672 |          7001.46 |           405.19 |               993.75 |\n",
        "| **M4**               |          791984 |        208148 |         35881.35 |           335.59 |               794.33 |\n",
        "| **M5**               |          698618 |        301514 |         35363.55 |           202.55 |               440.52 |\n",
        "| **M6**               |          756468 |        243664 |          4236.90 |           127.79 |               301.74 |\n",
        "| **UNet**             |          449578 |         51128 |        124987.51 |           610.96 |              3681.75 |\n",
        "| **T-RNN Dec**        |          405607 |         94459 |          4447.74 |           128.05 |               662.44 |\n",
        "| **T-RNN Enc**        |          410933 |         88880 |          5112.82 |           127.97 |               731.20 |"
      ],
      "metadata": {
        "id": "0cvfAbQ3B5bc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Listing the data for the studied application\n",
        "!gsutil ls gs://gresearch/prime"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "cellView": "form",
        "id": "1_BsXdRIKbq6",
        "outputId": "a11a3083-a571-434e-ecbf-11565a091cfc"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "gs://gresearch/prime/MobileNetEdgeTPU/\n",
            "gs://gresearch/prime/MobilenetEdgeTPU/\n",
            "gs://gresearch/prime/MobilenetV2/\n",
            "gs://gresearch/prime/MobilenetV3/\n",
            "gs://gresearch/prime/m4/\n",
            "gs://gresearch/prime/m5/\n",
            "gs://gresearch/prime/m6/\n",
            "gs://gresearch/prime/t_rnn_dec/\n",
            "gs://gresearch/prime/t_rnn_enc/\n",
            "gs://gresearch/prime/u-net/\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Importing the necessary libraries\n",
        "import tensorflow as tf\n",
        "import numpy as np"
      ],
      "metadata": {
        "id": "NsrEON0DKvIL"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title APIs for parsing PRIME datasets\n",
        "def parse_prime_tfrecords(proto):\n",
        "  prime_feature_description = {  \n",
        "    'param_1': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'param_2': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'param_3': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'param_4': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'param_5': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'param_6': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'param_7': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'param_8': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'param_9': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'param_10': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'runtime': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'area': tf.io.FixedLenFeature([], tf.float32),\n",
        "    'infeasible':tf.io.FixedLenFeature([], tf.int64),\n",
        "  }\n",
        "  return tf.io.parse_single_example(proto, prime_feature_description)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "uMWHVyBHKpja"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Parsing the dataset for the studied application\n",
        "model_name = 'm4' #@param [\"MobilenetEdgeTPU\", \"MobilenetV2\", \"MobilenetV3\", \"m4\", \"m5\", \"m6\", \"t_rnn_dec\", \"t_rnn_enc\", \"u-net\"]\n",
        "filenames = tf.io.gfile.glob(f'gs://gresearch/prime/{model_name}/*.tfrecord')\n",
        "raw_dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=64)\n",
        "parsed_dataset = raw_dataset.map(parse_prime_tfrecords)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "WAZFGWMgKuFR"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Reproducing the data in the Table 1\n",
        "number_of_infeasibles = 0\n",
        "number_of_feasibles = 0\n",
        "latency = []\n",
        "# Parsing the whole dataset for an application may take 15-20 minutes.\n",
        "for p in parsed_dataset:\n",
        "  if p['infeasible'] == 0:\n",
        "    number_of_feasibles+=1\n",
        "    latency.append(p['runtime'])\n",
        "  else:\n",
        "    number_of_infeasibles+=1\n",
        "print(f'{number_of_infeasibles},{number_of_feasibles},{max(latency)},{min(latency)},{np.average(latency)}')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rAYVi6wqLVGb",
        "outputId": "60033c83-59a1-4512-8fc5-aef46628fa7d"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "791984,208148,35881.3515625,335.59100341796875,794.33349609375\n"
          ]
        }
      ]
    }
  ]
}